"""
pandasAI的原理和介绍：
https://bohrium.dp.tech/notebooks/8242832145
"""

import pandasai as pai
from pandasai_local import LocalLLM
from config import settings
from typing import Dict, List, Any, Union
from matplotlib import pyplot, font_manager
from matplotlib.pylab import mpl
from pandasai import Agent
from loguru import logger

# 模型设置
llm = LocalLLM(
    api_base=settings.LLM_URI,
    model=settings.LLM_MODEL,
    api_key=settings.LLM_API_KEY,
)
pai.config.set({"llm": llm, "verbose": False, "save_logs": False, "max_retries": 3, "direct_sql": False})

# 中文字体设置
fonts = [f.name for f in font_manager.fontManager.ttflist]
cn_fonts = ["SimHei", "Microsoft YaHei", "FangSong", "Times New Roman"]
cn_fonts = list(set(cn_fonts).intersection(fonts))
mpl.rcParams["font.sans-serif"] = cn_fonts
pyplot.rcParams['font.family'] = cn_fonts
pyplot.rc("axes", unicode_minus = False)


def create_pai_agent(datasets: Union[Dict[str, Any], List[Dict[str, Any]]]) -> Agent:
    """创建PandasAI智能体"""
    if isinstance(datasets, dict):
        datasets = [datasets]
    dfs = [pai.DataFrame(dataset['rows']) for dataset in datasets]
    description = '\n'.join([dataset['description'] for dataset in datasets])
    agent = pai.Agent(dfs, description=description)
    return agent


def analyze_task(agent: Agent, task_name: str, out_type: str, retry_times: int=3):
    """调用PandasAI分析数据, 加入重试机制. """
    while retry_times > 0:
        try:
            output = agent.chat(query=task_name, output_type=out_type)
            if str(output.value).find('Please try again') == -1:
                return output
        except Exception as e:
            logger.error(e)
        retry_times -= 1
    return None
